using System;
using System.Collections.Generic;
using System.Linq;
using ArmyManager.Units;

namespace ArmyManager.Strategies
{
	public class BaseVisitorStrategy : IVisitorStrategy
	{
		private readonly Unit.UnitTypes[] _unitOrdering;
		private readonly bool _recurseFirst;

		public BaseVisitorStrategy(Unit.UnitTypes[] unitOrdering, bool recurseFirst)
		{
			_unitOrdering = unitOrdering;
			_recurseFirst = recurseFirst;
		}

		public IEnumerable<Unit> Initialize(IEnumerable<Unit> units)
		{
			UnitComparer unitComparer = new UnitComparer(_unitOrdering);
			Func<IEnumerable<Unit>, UnitComparer, IEnumerable<Unit>> enumerateMethod;
			if (_recurseFirst)
			{
				enumerateMethod = GetSubUnits;
			}
			else
			{
				enumerateMethod = GetNextUnits;
			}
			foreach (Unit unit in enumerateMethod(units, unitComparer))
			{
				yield return unit;
			}
		}

		private static IEnumerable<Unit> GetSubUnits(IEnumerable<Unit> currentUnits, IComparer<Unit.UnitTypes> unitComparer)
		{
			foreach (Unit currentUnit in currentUnits.OrderByDescending(u => u.UnitType, unitComparer))
			{
				yield return currentUnit;
				foreach (Unit subUnit in GetSubUnits(currentUnit.SubUnits, unitComparer))
				{
					yield return subUnit;
				}
			}
		}

		private static IEnumerable<Unit> GetNextUnits(IEnumerable<Unit> currentUnits, IComparer<Unit.UnitTypes> unitComparer)
		{
			List<Unit> nextUnits = new List<Unit>();
			foreach (Unit currentUnit in currentUnits.OrderByDescending(u => u.UnitType, unitComparer))
			{
				yield return currentUnit;
				nextUnits.AddRange(currentUnit.SubUnits);
			}
			if (nextUnits.Count > 0)
			{
				foreach (Unit nextUnit in GetNextUnits(nextUnits, unitComparer))
				{
					yield return nextUnit;
				}
			}
		}
	}
}